import logging
from torch.utils.data import DataLoader

import config
import data_loader

from model.Conv_TasNet.hand.conv_tasnet import ConvTasNet as Model
from train import ConvTasNetTrainV1

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def main():
    modelTrain = ConvTasNetTrainV1(device=config.device, retrain=False)

    logger.info(f"loaded model: {modelTrain.model.__class__.__name__}")
    x_train, y_train, x_cv, y_cv = data_loader.getXY(
        config.tr_directory, config.cv_directory
    )
    train_dataloader = DataLoader(
        data_loader.AudioData(x_train, y_train, config.device),
        batch_size=modelTrain.batch_size,
    )
    test_dataloader = DataLoader(
        data_loader.AudioData(x_cv, y_cv, config.device), modelTrain.batch_size
    )
    for X, y in test_dataloader:
        logger.info(f"Shape of X [N, C, H, W]: {X.shape}")
        logger.info(f"Shape of y: {y.shape} {y.dtype}")
        break
    modelTrain.train(train_dataloader, test_dataloader, modelTrain.output)
    pass


if __name__ == "__main__":
    main()
    pass
